import os

from fastapi import FastAPI, UploadFile, File, Form
from pydantic import BaseModel, Field

from api.PPOcrApi import ocr_image_cls, ocr_image_draw, ocr_image_mosaic_id, ocr_image_info, read_file
from api.utils import base64_to_ndarray, ndarray_to_base64

app = FastAPI(root_path="/api")


class OcrModel(BaseModel):
    img: str = Field(..., description="图片字符串的base64编码", min_length=50)


class QuestionModel(BaseModel):
    question: str = Field(..., description="提示词")


@app.post("/classify")
async def create_item(model: OcrModel):
    """
    图片分类接口
    :return:
    """
    img = base64_to_ndarray(model.img)
    res = ocr_image_cls(img)
    print(f"the result is {res}")
    return res


@app.post("/drawbox")
async def create_item(model: OcrModel):
    """
    图片Ocr drawbox接口
    :param model:
    :return:
    """
    img = base64_to_ndarray(model.img)
    image = ocr_image_draw(img)
    image_b64 = ndarray_to_base64(image)
    return {"image": image_b64}


@app.post("/mosaic/id")
async def create_item(model: OcrModel):
    """
    身份证号Ocr 打马赛克接口
    :param model:
    :return:
    """
    img = base64_to_ndarray(model.img)
    image = ocr_image_mosaic_id(img)
    image_b64 = ndarray_to_base64(image)
    return {"image": image_b64}


@app.post("/image/info")
async def create_item(model: OcrModel):
    """
    身份证号 、银行卡、营业执照 信息提取接口
    :param model:
    :return:
    """
    img = base64_to_ndarray(model.img)
    try:
        info = ocr_image_info(img)
        return {"info": info}
    except Exception as e:
        return {"info": str(e)}


@app.post('/chat')
async def create_item(question: str = Form(..., description="提示词"),
                      file: UploadFile = File(..., description="文件")):
    """
    上传 .txt、.docx、pdf文件,并对文档内容提问
    :param question:
    :param file:
    :return: gpt answer
    """
    try:
        content = read_file(file.filename, file.file, question)
        return {"answer": content}
    except Exception as e:
        return {"answer": str(e)}


if __name__ == "__main__":
    # 此参数必须设置
    # qwen or deepseek
    os.environ.setdefault("OPENAI_API_PLATE", "deepseek")

    # 默认使用deepseek api
    # 如何导入qwen 需要设置 api 地址  和 模型
    # os.environ.setdefault("OPENAI_API_URL", "http://localhost:11434/v1/")
    # 默认使用 qwen2.5:7b 、qwen2.5:14b、qwen2.5:32b
    # os.environ.setdefault("OPENAI_API_MODEL", "qwen2.5:7b")

    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8531)
